import datetime
import os
import logging
import hjson
import hvac
import pandas as pd

from absl import app, flags
from coin.support.accounting.logic.key_util import get_key_file
from coin.support.key_tool.vault.coin_key_high_dumper import (
    KeyInfo,
    list_keys as list_vault_keys)
from xunkemgmt_client.tool.slack_noti import send_to_slack
from xunkemgmt_client.client.util.query_util import query_accounts
from coin.proto.coin_request_pb2 import AccountInfoProto

FLAGS = flags.FLAGS

TRADE_KEY_EXPIERY_MAP = {
  'bybit': 90,
  'huobi': 90,
  'mexc': 90,
  'upbit': 365
}

def list_all_keys(key_root):
  key_infos = []
  for root, dirs, files in os.walk(key_root):
    for file in files:
      key_name, extension = os.path.splitext(file)
      if extension == '.json':
        key_path = os.path.join(root, file)
        with open(key_path) as fp:
          key_msg = hjson.load(fp)
        key_infos.append(KeyInfo(key_path, key_msg))
  return key_infos


def select_key_path(key_infos, exchanges):
  selected_key_infos = []
  exchange_per_key = []
  for key_info in key_infos:
    owner, exchange_name, key_name = key_info.key_path.split('/')[-3:]
    for exchange in exchanges:
      if exchange in exchange_name.lower():
        selected_key_infos.append(key_info)
        exchange_per_key.append(exchange)
        break
  assert len(selected_key_infos) == len(exchange_per_key)
  return selected_key_infos, exchange_per_key


def validate_key_msg(key_msg):
  assert 'generated' in key_msg and key_msg['generated'] is not None, key_msg
  assert 'key_name' in key_msg, key_msg


def check_key_expiry(key_infos, expiry_map, alarm_days):
  current_dt = datetime.date.today()
  key_infos, exchange_per_key = select_key_path(key_infos, list(expiry_map.keys()))
  alarm_msg = []
  for key_info, exchange in zip(key_infos, exchange_per_key):
    key_msg = key_info.key_content
    validate_key_msg(key_msg)
    # loan key is used by Huobi to monitor our assets and never gets expired.
    if exchange == 'huobi' and key_msg['key_name'] == 'loan':
      continue
    try:
      gen_dt = datetime.datetime.strptime(key_msg['generated'], '%Y%m%d').date()
    except ValueError:
      gen_dt = datetime.datetime.strptime(key_msg['generated'], '%Y-%m-%d').date()
    expiry_dt = gen_dt + datetime.timedelta(days=expiry_map[exchange])
    if expiry_dt - current_dt <= datetime.timedelta(days=alarm_days):
      alarm_msg.append([exchange, key_msg['owner'], expiry_dt, key_info.key_path])
  alarm_df = pd.DataFrame(alarm_msg, columns=['exchange', 'owner', 'expiry_date', 'key_file'])
  alarm_df = alarm_df.set_index(['exchange', 'owner'])
  alarm_df = alarm_df.sort_index()
  return alarm_df


def check_key_file_path(key_infos, key_type):
  key_paths = [key_info.key_path for key_info in key_infos]
  accounts = query_accounts(active=True, as_proto=True)
  if key_type == 'trade':
    accounts = [account for account in accounts 
    if account.market_type in ('Spot', 'Futures', 'Options')
    and account.account_type != AccountInfoProto.MASTER]
  std_key_paths = {}
  for account in accounts:
    try:
      assert account.exchange is not None
      assert account.api_version is not None
      assert account.owner is not None
      std_key_path = get_key_file(account.market_type,
                                  account.exchange,
                                  account.owner,
                                  account.api_version,
                                  key_type,
                                  check_exist=False)
      std_key_paths[std_key_path] = account
    except KeyError as e:
      logging.info('exchange key map not registered %s' % e)
  std_key_paths = {key_path: std_key_paths[key_path]
    for key_path in std_key_paths if key_path not in key_paths}
  exclude_list = []
  for std_key_path in std_key_paths:
    account_info = std_key_paths[std_key_path]
    exchange = std_key_path.split('/')[-2]
    for key_path in key_paths:
      owner, exchange_name, key_name = key_path.split('/')[-3:]
      if account_info.owner == owner and exchange == exchange_name:
        exclude_list.append((account_info.market_type,
                             account_info.exchange,
                             account_info.api_version,
                             account_info.owner,
                             key_path,
                             std_key_path))
  exclude_df = pd.DataFrame(list(set(exclude_list)),
    columns=['market_type', 'exchange', 'api_version',
             'owner', 'curr_key_path', 'std_key_path'])
  exclude_df = exclude_df.sort_values(
    ['market_type', 'exchange', 'api_version', 'owner'])
  return exclude_df


def main(_):
  pd.set_option('display.width', 200)
  pd.set_option('max_colwidth', 200)
  vault_addr = FLAGS.vault_addr
  key_root = FLAGS.key_root
  assert vault_addr is not None or key_root is not None
  if vault_addr:
    vault_token_file = FLAGS.vault_token_file
    assert os.path.exists(vault_token_file), vault_token_file
    with open(vault_token_file) as fp:
      vault_token = hjson.load(fp)['token']
    client = hvac.Client(url=vault_addr, token=vault_token)
    assert client.is_authenticated()
    key_infos = list_vault_keys(client, FLAGS.vault_path, FLAGS.vault_mount_point)
    key_path_naming_df = check_key_file_path(key_infos, 'trade')
  else:
    key_infos = list_all_keys(key_root)
    key_path_naming_df = check_key_file_path(key_infos, 'view')
  expiry_trade_key_df = check_key_expiry(
      key_infos, TRADE_KEY_EXPIERY_MAP, FLAGS.pre_alarm_days)
  print(expiry_trade_key_df.to_string())
  if FLAGS.slack_receiver:
    mention_list = FLAGS.mention_member.split(',') if FLAGS.mention_member else None
    if not expiry_trade_key_df.empty:
      msg = "```" + "key expiry alarm\n" + expiry_trade_key_df.to_string() + "```"
      send_to_slack(msg, FLAGS.slack_receiver, 'msg', mention_list=mention_list)
    if not key_path_naming_df.empty:
      msg = "```" + "please follow the key path naming rules:\n"
      msg += key_path_naming_df.to_string(index=False) + "```"
      send_to_slack(msg, FLAGS.slack_receiver, 'msg', mention_list=mention_list)
  return


if __name__ == '__main__':
  flags.DEFINE_string('vault_addr', None, '')
  flags.DEFINE_string('vault_token_file', None, '')
  flags.DEFINE_string('vault_path', '', '')
  flags.DEFINE_string('vault_mount_point', None, '')
  flags.DEFINE_integer('pre_alarm_days', 3, 'ahead days')
  flags.DEFINE_string('slack_receiver', None, 'slack channel')
  flags.DEFINE_string('mention_member', None, 'mention in slack')
  app.run(main)
